import heapq
from collections import defaultdict
from itertools import product
import matplotlib.pyplot as plt
import numpy as np



def space_time_a_star(start_positions, goal_positions, map, num_agents):
    # 将起始位置转换为元组
    start_positions = tuple(start_positions)
    goal_positions = tuple(goal_positions)
    
    # 初始化
    open_list = []
    closed_set = set()
    g_scores = defaultdict(lambda: float('inf'))
    f_scores = defaultdict(lambda: float('inf'))
    came_from = {}  # 用于存储路径信息
    
    # 初始状态
    initial_state = (start_positions, 0)  # (positions, time_step)
    g_scores[initial_state] = 0
    f_scores[initial_state] = heuristic(start_positions, goal_positions)
    
    # 将初始状态放入开放列表
    heapq.heappush(open_list, (f_scores[initial_state], initial_state))
    
    while open_list:
        _, current_state = heapq.heappop(open_list)
        current_positions, current_time = current_state
        
        if current_positions == goal_positions:
            return reconstruct_path(came_from, current_state)
        
        closed_set.add(current_state)
        
        for neighbor in get_neighbors(current_positions, current_time, map, num_agents,goal_positions):
            neighbor_positions, neighbor_time = neighbor
            
            if neighbor in closed_set:
                continue
            
            tentative_g_score = g_scores[current_state] + 1  # 移动一步的成本为1
            
            if tentative_g_score < g_scores[neighbor]:
                came_from[neighbor] = current_state  # 记录父节点
                g_scores[neighbor] = tentative_g_score
                f_scores[neighbor] = tentative_g_score + heuristic(neighbor_positions, goal_positions)
                heapq.heappush(open_list, (f_scores[neighbor], neighbor))
    
    return None  # 未找到解

def heuristic(positions, goals):
    # 返回启发式值，例如曼哈顿距离之和
    return sum(abs(p[0] - g[0]) + abs(p[1] - g[1]) for p, g in zip(positions, goals))

def get_neighbors(positions, current_time, map, num_agents, goal_positions):
    # 生成每个智能体的所有可能移动，并根据启发式值排序
    individual_moves = []
    for i in range(num_agents):
        moves_for_agent = []
        for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            new_x, new_y = positions[i][0] + dx, positions[i][1] + dy
            if is_valid(new_x, new_y, map):
                h_value = abs(new_x - goal_positions[i][0]) + abs(new_y - goal_positions[i][1])
                moves_for_agent.append(((new_x, new_y), h_value))
        # 根据启发式值对移动进行排序
        moves_for_agent.sort(key=lambda x: x[1])
        individual_moves.append([move[0] for move in moves_for_agent])

    # 动态生成有效的邻居状态
    valid_combinations = []
    for combination in product(*individual_moves):
        if not has_conflict(combination, current_time) and not is_direct_swap(positions, combination):
            valid_combinations.append((tuple(combination), current_time + 1))

    # 可以进一步根据整体启发式值对有效组合进行排序
    valid_combinations.sort(key=lambda x: sum(abs(p[0] - g[0]) + abs(p[1] - g[1]) for p, g in zip(x[0], goal_positions)))

    return valid_combinations

def is_direct_swap(old_positions, new_positions):
    # 检查是否有两个智能体直接交换位置
    for i in range(len(old_positions)):
        for j in range(i + 1, len(old_positions)):
            if old_positions[i] == new_positions[j] and old_positions[j] == new_positions[i]:
                return True
    return False

def has_conflict(positions, current_time):
    # 检查新位置是否与其他智能体发生冲突
    seen = set()
    for pos in positions:
        if pos in seen:
            return True
        seen.add(pos)
    return False

def is_valid(x, y, map):
    # 检查位置 (x, y) 是否在地图范围内且不是障碍物
    return 0 <= x < len(map) and 0 <= y < len(map[0]) and map[x][y] != 1


def reconstruct_path(came_from, final_state):
    # 重构路径
    path = []
    current = final_state
    while current in came_from:
        path.append(current)
        current = came_from[current]
    path.append(current)
    path.reverse()  # 反转路径以获得从开始到结束的顺序
    return path

import heapq
from collections import defaultdict
from itertools import product
import matplotlib.pyplot as plt
import numpy as np

# 之前的函数定义保持不变...

def visualize_path(path, map):
    # 获取地图尺寸
    height = len(map)
    width = len(map[0])
    
    # 创建一个新的图形
    fig, ax = plt.subplots()
    
    # 绘制栅格
    for x in range(width + 1):
        ax.axvline(x, color='black', linewidth=1)
    for y in range(height + 1):
        ax.axhline(y, color='black', linewidth=1)
    
    # 设置坐标轴范围
    ax.set_xlim(0, width)
    ax.set_ylim(0, height)
    
    # 反转y轴，使得原点在左下角
    ax.invert_yaxis()
    
    # 移除坐标轴标签
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    
    # 显示障碍物
    for x in range(width):
        for y in range(height):
            if map[y][x] == 1:
                rect = plt.Rectangle((x, y), 2, 1, facecolor='gray')
                ax.add_patch(rect)
    
    # 动态更新智能体位置
    for state in path:
        positions, time = state
        ax.cla()  # 清除当前图形
        
        # 重新绘制栅格和障碍物
        for x in range(width + 1):
            ax.axvline(x, color='black', linewidth=1)
        for y in range(height + 1):
            ax.axhline(y, color='black', linewidth=1)
        
        for x in range(width):
            for y in range(height):
                if map[y][x] == 1:
                    rect = plt.Rectangle((x, y), 1, 1, facecolor='gray')
                    ax.add_patch(rect)
        
        # 绘制智能体位置
        colors = ['red', 'blue', 'green', 'orange', 'purple']  # 每个智能体一种颜色
        for i, (x, y) in enumerate(positions):
            circle = plt.Circle((x + 0.5, y + 0.5), 0.3, color=colors[i % len(colors)])
            ax.add_patch(circle)
        
        # 更新标题以显示当前时间步
        ax.set_title(f'Time {time}')
        
        # 显示并暂停一段时间
        plt.pause(1)
    
    plt.show()


if __name__ == '__main__':
    # 示例用例
    map = [
        [0, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 0]
    ]

    start_positions = [(0, 0), (3, 3), (3, 2), (2, 3), (2, 2)]
    goal_positions = [(3, 3), (0, 0), (0, 1), (2, 1), (0, 2)]
    num_agents = 5

    result = space_time_a_star(start_positions, goal_positions, map, num_agents)
    if result:
        print("Path found:")
        for state in result:
            positions, time = state
            print(f"Time {time}: {positions}")
        
        # 可视化路径
        visualize_path(result, map)
    else:
        print("No path found.")